import requests
import jsonpath_rw
import json
import traceback
import re

'''
提取器
'''


class Extract:
    pb, pe = '\$\{', '\}'  # \为转义   ${}

    def __init__(self, response: requests.models.Response = None):
        self.response = response
        self.pub_compile = re.compile(r'{}(.+?){}'.format(self.pb, self.pe))

    def get_pub_param(self, string) -> list:
        return list(set(self.pub_compile.findall(string)))

    def set_pub_param(self, old_v, new_v, string):
        return re.sub(r'{}{}{}'.format(self.pb, old_v, self.pe), new_v, string)

    def status_code(self, *args):
        return self.response.status_code

    def headers(self, expr):
        return self.response.headers.get(expr)

    def json(self, expr):
        result = None
        json_str = self.response.text
        try:
            json_obj = json.loads(json_str)
        except:
            # json解析失败
            raise Exception('响应体不符合JSON类型')
        try:
            jsonpath_expr = jsonpath_rw.parse(expr)
        except:
            raise Exception("提取表达式格式错误,'{}'".format(expr))

        re_obj = jsonpath_expr.find(json_obj)  # todo 返回所有复合条件的结果，当前仅取第一条
        result = re_obj[0].value if re_obj else None

        return result


if __name__ == '__main__':
    e = Extract()
    abc, sdf = '|qwerty|', '|asq23e|'
    str = 'asd${abc}89kmasdk${sdf}'
    li = (['abc', abc], ['sdf', sdf])
    for item in li:
        str = e.set_pub_param(item[0], item[1], str)
    print(str)
